from tqdm import tqdm
import numpy as np
import torch
import collections
import random

from buffer.Buffer import Buffer


class ReplayBuffer(Buffer):
    """ 经验回放池 """

    def __init__(self, capacity):
        # 队列
        super().__init__()
        self.buffer = collections.deque(maxlen=capacity)

    def add(self, state, action, reward, next_state, done):  # 将数据存入buffer中
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):  # 从buffer中采样数据，数据量为batch_size
        transitions = random.sample(self.buffer, batch_size)
        state, action, reward, next_state, done = zip(*transitions)
        return np.array(state), np.array(action), np.array(reward), np.array(next_state), np.array(done)

    def size(self):  # 目前buffer中数据的数量
        return len(self.buffer)
